#include "kernel_operator.h"
#include "is_real_obj_custom_tiling.h"
#include <cstdint>
#include <cmath>

constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue

constexpr uint32_t ISREALOBJ_BFLOAT16 = 0;            // 对应 DT_BF16，bf16 类型
constexpr uint32_t ISREALOBJ_FLOAT16 = 1;             // 对应 DT_FLOAT16，half（float16）类型
constexpr uint32_t ISREALOBJ_FLOAT32 = 2;             // 对应 DT_FLOAT，float32 类型
constexpr uint32_t ISREALOBJ_INT8 = 3;                // 对应 DT_INT8，int8 类型
constexpr uint32_t ISREALOBJ_INT16 = 4;               // 对应 DT_INT16，int16 类型
constexpr uint32_t ISREALOBJ_INT32 = 5;               // 对应 DT_INT32，int32 类型
constexpr uint32_t ISREALOBJ_UINT8 = 6;               // 对应 DT_UINT8，u8（无符号int8）类型
constexpr uint32_t ISREALOBJ_UINT16 = 7;              // 对应 DT_UINT16，u16（无符号int16）类型
constexpr uint32_t ISREALOBJ_UINT32 = 8;              // 对应 DT_UINT32，u32（无符号int32）类型
constexpr uint32_t ISREALOBJ_INT64 = 9;               // 对应 DT_INT64，int64 类型
constexpr uint32_t ISREALOBJ_UINT64 = 10;             // 对应 DT_UINT64，u64（无符号int64）类型
constexpr uint32_t ISREALOBJ_BOOL = 11;               // 对应 DT_BOOL，bool 类型
constexpr uint32_t ISREALOBJ_DOUBLE = 12;             // 对应 DT_DOUBLE，double（float64）类型
constexpr uint32_t ISREALOBJ_STRING = 13;             // 对应 DT_STRING，string（字符串）类型
constexpr uint32_t ISREALOBJ_DUAL_SUB_INT8 = 14;      // 对应 DT_DUAL_SUB_INT8，dual output int8（双输出int8）类型
constexpr uint32_t ISREALOBJ_DUAL_SUB_UINT8 = 15;     // 对应 DT_DUAL_SUB_UINT8，dual output u8（双输出无符号int8）类型
constexpr uint32_t ISREALOBJ_COMPLEX64 = 16;          // 对应 DT_COMPLEX64，complex64（复数64位）类型
constexpr uint32_t ISREALOBJ_COMPLEX128 = 17;         // 对应 DT_COMPLEX128，complex128（复数128位）类型
constexpr uint32_t ISREALOBJ_QINT8 = 18;              // 对应 DT_QINT8，qint8（量化int8）类型
constexpr uint32_t ISREALOBJ_QINT16 = 19;             // 对应 DT_QINT16，qint16（量化int16）类型
constexpr uint32_t ISREALOBJ_QINT32 = 20;             // 对应 DT_QINT32，qint32（量化int32）类型
constexpr uint32_t ISREALOBJ_QUINT8 = 21;             // 对应 DT_QUINT8，quint8（量化无符号int8）类型
constexpr uint32_t ISREALOBJ_QUINT16 = 22;            // 对应 DT_QUINT16，quint16（量化无符号int16）类型
constexpr uint32_t ISREALOBJ_RESOURCE = 23;           // 对应 DT_RESOURCE，resource（资源）类型
constexpr uint32_t ISREALOBJ_STRING_REF = 24;         // 对应 DT_STRING_REF，string ref（字符串引用）类型
constexpr uint32_t ISREALOBJ_DUAL = 25;               // 对应 DT_DUAL，dual output（双输出）类型
constexpr uint32_t ISREALOBJ_VARIANT = 26;            // 对应 DT_VARIANT，dt_variant（变体）类型
constexpr uint32_t ISREALOBJ_UNDEFINED = 27;          // 对应 DT_UNDEFINED，未设置的DataType字段
constexpr uint32_t ISREALOBJ_INT4 = 28;               // 对应 DT_INT4，int4 类型
constexpr uint32_t ISREALOBJ_UINT1 = 29;              // 对应 DT_UINT1，u1（无符号int1）类型
constexpr uint32_t ISREALOBJ_INT2 = 30;               // 对应 DT_INT2，int2 类型
constexpr uint32_t ISREALOBJ_UINT2 = 31;              // 对应 DT_UINT2，u2（无符号int2）类型
constexpr uint32_t ISREALOBJ_COMPLEX32 = 32;          // 对应 DT_COMPLEX32，complex32（复数32位）类型
constexpr uint32_t ISREALOBJ_HIFLOAT8 = 33;           // 对应 DT_HIFLOAT8，hif8（高精度float8）类型
constexpr uint32_t ISREALOBJ_FLOAT8_E5M2 = 34;        // 对应 DT_FLOAT8_E5M2，fp8_e5m2（float8格式，5位指数2位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT8_E4M3FN = 35;      // 对应 DT_FLOAT8_E4M3FN，fp8_e4m3（float8格式，4位指数3位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT8_E8M0 = 36;        // 对应 DT_FLOAT8_E8M0，fp8_e8m0（float8格式，8位指数0位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT6_E3M2 = 37;        // 对应 DT_FLOAT6_E3M2，fp6_e3m2（float6格式，3位指数2位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT6_E2M3 = 38;        // 对应 DT_FLOAT6_E2M3，fp6_e2m3（float6格式，2位指数3位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT4_E2M1 = 39;        // 对应 DT_FLOAT4_E2M1，fp4_e2m1（float4格式，2位指数1位尾数）类型
constexpr uint32_t ISREALOBJ_FLOAT4_E1M2 = 40;        // 对应 DT_FLOAT4_E1M2，fp4_e1m2（float4格式，1位指数2位尾数）类型
constexpr uint32_t ISREALOBJ_MAX = 41;                // 对应 DT_MAX，标记AscendC数据类型的边界

constexpr uint32_t LAST_TWO_TILE = 2;

// 定义复数类型
struct complex64 {
    float real;
    float imag;
};

// 定义复数类型
struct complex128 {
    float real_1;
    float real_2;
    float imag_1;
    float imag_2;
};

template <typename dataType>
class KernelIsRealObj {
public:
    __aicore__ inline KernelIsRealObj() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR out, IsRealObjCustomTilingData tiling)
    {
        this->dataTypeId = tiling.dataType;
        if (tiling.isEvenCore)
        {
            this->blockLength = tiling.blockLength;
            this->tileNum = tiling.tileNum;
            this->tileLength = tiling.tileLength / BUFFER_NUM;
            this->lastTileLength = tiling.lastTileLength;

            xGm.SetGlobalBuffer((__gm__ dataType *)x + this->blockLength * AscendC::GetBlockIdx(), this->blockLength);
            outGm.SetGlobalBuffer((__gm__ int8_t *)out, 1);  // 输出只需要一个元素
        }
        else
        {
            if (AscendC::GetBlockIdx() < tiling.formerNum)
            {
                this->tileNum = tiling.formerTileNum;
                this->tileLength = tiling.formerTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.formerLastTileLength;

                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * AscendC::GetBlockIdx(),
                                    tiling.formerLength);
                outGm.SetGlobalBuffer((__gm__ int8_t *)out, 1);  // 输出只需要一个元素
            }
            else
            {
                this->tileNum = tiling.tailTileNum;
                this->tileLength = tiling.tailTileLength / BUFFER_NUM;
                this->lastTileLength = tiling.tailLastTileLength;

                xGm.SetGlobalBuffer((__gm__ dataType *)x + tiling.formerLength * tiling.formerNum +
                                        tiling.tailLength * (AscendC::GetBlockIdx() - tiling.formerNum),
                                    tiling.tailLength);
                outGm.SetGlobalBuffer((__gm__ int8_t *)out, 1);  // 输出只需要一个元素
            }
        }
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(dataType));
        pipe.InitBuffer(outQueueBool, BUFFER_NUM, 1);
    }

    __aicore__ inline void Process() {
        CopyIn(0);
        Compute(0);
        CopyOut(0);
    }

private:
    __aicore__ inline void CopyIn(int32_t progress) {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.AllocTensor<dataType>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1)))
        {
            AscendC::DataCopy(xLocal, xGm[(progress - LAST_TWO_TILE) * this->tileLength + this->lastTileLength],
                              this->tileLength);
        }
        else
        {
            AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
        }
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress) {
        AscendC::LocalTensor<dataType> xLocal = inQueueX.DeQue<dataType>();
        AscendC::LocalTensor<int8_t> outLocal = outQueueBool.AllocTensor<int8_t>();
        
        bool isComplex = (dataTypeId == ISREALOBJ_COMPLEX64 || dataTypeId == ISREALOBJ_COMPLEX128);

        if (!isComplex) {
            outLocal.SetValue(0, 1);
        }
        else{
            outLocal.SetValue(0, 0);
        }

        outQueueBool.EnQue(outLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress) {
        AscendC::LocalTensor<int8_t> outLocal = outQueueBool.DeQue<int8_t>();
        AscendC::DataCopyExtParams copyParams{1, (uint32_t)(1), 0, 0, 0};

        if((progress == (this->tileNum * BUFFER_NUM - 2)) || (progress == (this->tileNum * BUFFER_NUM - 1))){
            AscendC::DataCopyPad(outGm[0], outLocal, copyParams);
        }else{
            AscendC::DataCopyPad(outGm[0], outLocal, copyParams);
        }
        outQueueBool.FreeTensor(outLocal);
    }
    
private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueBool;
    AscendC::GlobalTensor<dataType> xGm;
    AscendC::GlobalTensor<int8_t> outGm;
    AscendC::DataCopyExtParams copyParams;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;

    uint32_t dataTypeId;
};


extern "C" __global__ __aicore__ void is_real_obj_custom(GM_ADDR x, GM_ADDR out, IsRealObjCustomTilingData tiling)
{
    if (tiling.dataType == ISREALOBJ_BFLOAT16) {
        KernelIsRealObj<bfloat16_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == ISREALOBJ_FLOAT16) {
        KernelIsRealObj<half> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == ISREALOBJ_FLOAT32) {
        KernelIsRealObj<float> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == ISREALOBJ_INT8) {
        KernelIsRealObj<int8_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == ISREALOBJ_INT16) {
        KernelIsRealObj<int16_t> op;
        op.Init(x, out, tiling);
        op.Process();
    } else if (tiling.dataType == ISREALOBJ_INT32) {
        KernelIsRealObj<int32_t> op;
        op.Init(x, out, tiling);
        op.Process();
    }else if (tiling.dataType == ISREALOBJ_COMPLEX64) {
        KernelIsRealObj<complex64> op;
        op.Init(x, out, tiling);
        op.Process();
    }else if (tiling.dataType == ISREALOBJ_COMPLEX128) {
        KernelIsRealObj<complex128> op;
        op.Init(x, out, tiling);
        op.Process();
    } 
    else {
        return;
    }
}
